/*!
  @file stereo_gc.cpp
  @copyright 2013 Kubota Lab. All rights resereved.
*/

#include <limits>

#include <boost/format.hpp>

#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/highgui/highgui.hpp>

#define MODULE_NAME "StereoGC"
#define DEBUG_FLAG  debug_
#include "kcv/common/io.h"

#include "kcv/stereo/stereo_gc.h"


namespace kcv
{

void StereoGC::process(const cv::Mat &imL, const cv::Mat &imR, cv::Mat &disparity) 
{ 
  /* init */
  initImages(imL, imR);

  cv::Ptr<cv::CLAHE> clahe = cv::createCLAHE();
  clahe->apply(imL_, imL_);
  clahe->apply(imR_, imR_);

  show_image();

  clock_t start = clock();
  performAlphaExtension();
  clock_t end = clock();
  INFO("Time(Stereo): " << (boost::format("%.2f") %((double)(end-start)/CLOCKS_PER_SEC)).str() << " [s]\n");


  /* copy disparity maps */
  disparity = dis_.clone();
}

void StereoGC::initImages(const cv::Mat &imL, const cv::Mat &imR)
{
  /* color conversion */
  cv::Mat imL_gray(imL.size(), CV_8U);
  cv::Mat imR_gray(imR.size(), CV_8U);
  if (imL.channels() == 3)
  {
    cv::cvtColor(imL, imL_gray, CV_BGR2GRAY);
    cv::cvtColor(imR, imR_gray, CV_BGR2GRAY);
  }
  else
  {
    imL_gray = imL;
    imR_gray = imR;
  }

  /* scale conversion */
  if (scale_ == 1.0)
  {
    imL_ = imL_gray;
    imR_ = imR_gray;
  }
  else
  {
    cv::Size sz(imL.cols * scale_, imL.rows * scale_);
    imL_.create(sz, CV_8U);
    imR_.create(sz, CV_8U);
    cv::resize(imL_gray, imL_, sz);
    cv::resize(imR_gray, imR_, sz);
  }

  /* create disparity buffers */
  dis_ = cv::Mat::zeros(imL_.size(), CV_8U);
}

void StereoGC::performAlphaExtension()
{
  int n = imL_.rows;
  int m = imL_.cols;
  DEBUG("Performing Alpha Extension for " << m << "x" << n << " images....");

  size_t num_site = n * m;
  size_t num_pair = (n - 1) * m + n * (m - 1);
  size_t num_node = num_site + num_pair;
  size_t num_edge = 2 * num_site + 4 * num_pair;
  GraphPtr g(new GraphType(num_node, num_edge));

  long E = std::numeric_limits<long>::max();
  for (size_t num_itr = 0; num_itr < 10000; ++num_itr)
  {
    bool updated = false;
    for (int a = max_disp_; a >= min_disp_; --a)
    {
      g->add_node(num_site);

      /* 
       * data term 
       */
      int lam_g_ = 1;
      for (int i = 0; i < n; ++i)
      {
        for (int j = 0; j < m; ++j)
        {
          int c_sv = (dis_.at<uchar>(i, j) == a)? 
                       std::numeric_limits<int>::max(): 
                       calcDataCost(i, j, dis_.at<uchar>(i, j));
          int c_vt = calcDataCost(i, j, a);
          g->add_tweights(SITEIDX(i, j), lam_g_ * c_sv, lam_g_ * c_vt);
        }
      }

      /* 
       * smoothing term 
       */
      int lam_h_ = 4;
      int idx = num_site;

      /* on epipolar line */
      for (int i = 0; i < n; ++i)
      {
        for (int j = 0; j < m - 1; ++j)
        {
          int u = dis_.at<uchar>(i, j);
          int v = dis_.at<uchar>(i, j + 1);
          if (u == v)
          {
            int c_uv = calcSmoothCost(u, a);
            g->add_edge(SITEIDX(i, j), SITEIDX(i, j + 1), lam_h_ * c_uv, lam_h_ * c_uv);
          }
          else
          {
            int c_suv = calcSmoothCost(u, v);
            int c_uuv = calcSmoothCost(u, a);
            int c_vuv = calcSmoothCost(v, a);

            g->add_node();
            g->add_tweights(idx, c_suv, 0);
            g->add_edge(SITEIDX(i, j),     idx, lam_h_ * c_uuv, lam_h_ * c_uuv);
            g->add_edge(SITEIDX(i, j + 1), idx, lam_h_ * c_vuv, lam_h_ * c_vuv);
            ++idx;
          }
        }
      }

      /* between adjuscent epipolar lines */
      for (int i = 0; i < n - 1; ++i)
      {
        for (int j = 0; j < m; ++j)
        {
          int u = dis_.at<uchar>(i, j);
          int v = dis_.at<uchar>(i + 1, j);
          if (u == v)
          {
            int c_uv = calcSmoothCost(u, a);
            g->add_edge(SITEIDX(i, j), SITEIDX(i + 1, j), lam_h_ * c_uv, lam_h_ * c_uv);
          }
          else
          {
            int c_suv = calcSmoothCost(u, v);
            int c_uuv = calcSmoothCost(u, a);
            int c_vuv = calcSmoothCost(a, v);

            g->add_node();
            g->add_tweights(idx, c_suv, 0);
            g->add_edge(SITEIDX(i, j),     idx, lam_h_ * c_uuv, lam_h_ * c_uuv);
            g->add_edge(SITEIDX(i + 1, j), idx, lam_h_ * c_vuv, lam_h_ * c_vuv);
            ++idx;
          }
        }
      }

      /* max-flow, min-cut */
      long E_new = g->maxflow();
      if (E_new < E)
      {
        E = E_new;
        updated = true;
          
        /* extend alpha */
        for (int i = 0; i < n; ++i)
        {
          for (int j = 0; j < m; ++j)
          {
            if (g->what_segment(SITEIDX(i, j)) == GraphType::SOURCE)
            {
              dis_.at<uchar>(i, j) = a;
            }
          }
        }
      }
      g->reset();

      if (debug_) show_disparity();
    } 

    DEBUG("New energy (" << num_itr << "): " << E)

    /* finish loop if there are no update */
    if (!updated) break;
  }
}

int StereoGC::calcDataCost(const int i, const int j, const int d)
{
  //return abs(imL_.at<uchar>(i, j) - imR_.at<uchar>(i, j - d));
  int cost = 0;
  switch (func_data_)
  {
    case DATA_COST_SAD:
      cost = sad(i, j, d);
      break;
    case DATA_COST_ZNCC:
      cost = zncc(i, j, d);
      break;
    case DATA_COST_DIFF:
      cost = diff(i, j, d);
      break;
    default:
      ERROR("Invalid data cost function");
      break;
  }
  return cost;
}

int StereoGC::sad(const int i, const int j, const int d)
{
  int hwin = win_size_ / 2;
  if (!isBetween(i - hwin, 0, imL_.rows)
      || !isBetween(i + hwin, 0, imL_.rows)
      || !isBetween(j - hwin - d, 0, imL_.cols)
      || !isBetween(j + hwin, 0, imL_.cols))
  {
    return LARGE_COST;
  }

  int sum = 0;
  for (int s = -hwin; s <= hwin; ++s)
  {
    for (int t = -hwin; t <= hwin; ++t)
    {
      if (isBetween(i + s, 0, imL_.rows)
          && isBetween(j + t, 0, imL_.cols)
          && isBetween(j + t - d, 0, imL_.cols))
      {
        sum += abs(imL_.at<uchar>(i + s, j + t) - imR_.at<uchar>(i + s, j + t - d));
      }
    }
  }
  return sum / (win_size_ * win_size_);
}

int StereoGC::zncc(const int i, const int j, const int d)
{
  int hwin = win_size_ / 2;
  if (!isBetween(i - hwin, 0, imL_.rows)
      || !isBetween(i + hwin, 0, imL_.rows)
      || !isBetween(j - hwin - d, 0, imL_.cols)
      || !isBetween(j + hwin, 0, imL_.cols))
  {
    return LARGE_COST;
  }

  double sumL = 0;
  double sqsumL = 0;
  double sumR = 0;
  double sqsumR = 0;
  double sqsumLR = 0;
  for (int s = -hwin; s <= hwin; ++s)
  {
    for (int t = -hwin; t <= hwin; ++t)
    {
      sumL += imL_.at<uchar>(i + s, j + t);
      sumR += imR_.at<uchar>(i + s, j + t - d);
      sqsumL += pow(imL_.at<uchar>(i + s, j + t), 2);
      sqsumR += pow(imR_.at<uchar>(i + s, j + t - d), 2);
      sqsumLR += imL_.at<uchar>(i + s, j + t) * imR_.at<uchar>(i + s, j + t - d);
    }
  }

  double N = win_size_ * win_size_;
  double eps = 0.00001;
  double zncc = (N * sqsumLR - sumL * sumR) / sqrt((N * sqsumL - sumL * sumL + eps) * (N * sqsumR - sumR * sumR + eps));
  //fprintf(stdout, "%.2f %.2f %.2f %.2f %.2f %.2f\n", sumL, sumR, sqsumL, sqsumR, sqsumLR, zncc);
  //INFO(zncc);
  return 30 * (1.0 - zncc);
}

int StereoGC::diff(const int i, const int j, const int d)
{
  return abs(imL_.at<uchar>(i, j) - imR_.at<uchar>(i, j - d));
}

int StereoGC::calcSmoothCost(const int u, const int v)
{
  //return (u == v)? 0: 1;  // Potts
  return abs(u - v);  // 
}

void StereoGC::show_image(const std::string &winname)
{
  cv::Mat disp(imL_.rows, 2 * imL_.cols, imL_.type());
  cv::Mat r = disp(cv::Rect(0, 0, imL_.cols, imL_.rows));
  imL_.copyTo(r);
  r = disp(cv::Rect(imL_.cols, 0, imL_.cols, imL_.rows));
  imR_.copyTo(r);

  for (int i = 0; i < imL_.rows; i += 10)
  {
    cv::line(disp, cv::Point(0, i), cv::Point(2 * imL_.cols - 1, i), 0, 1);
  }
  cv::imshow("Image", disp);
  cv::waitKey(1);
}

void StereoGC::show_disparity(const std::string &winname)
{
  cv::Mat disp = dis_.clone();
  disp *= 255. / max_disp_;
  cv::imshow("disparity", disp);
  cv::waitKey(1);
}

}  // namespace kcv

